package com.kaithan.framework.common.threadpool;

import lombok.extern.slf4j.Slf4j;

import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CompletionException;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.Executor;
import java.util.function.BiConsumer;
import java.util.function.BiFunction;
import java.util.function.Function;
import java.util.function.Supplier;

/**
 * 类描述：封装 CompletableFuture 类，使得线程池执行过程中出异常能够被明显觉察到
 *
 * @author ZhuYin
 * @since 2024年12月27日
 */
@Slf4j
public class MyCompletableFuture<T> extends CompletableFuture<T> {

    // 全局异常处理器
    private static BiConsumer<? super Throwable, ? super String> globalExceptionHandler;

    /**
     * 设置全局异常处理器
     */
    public static void setGlobalExceptionHandler(BiConsumer<? super Throwable, ? super String> handler) {
        MyCompletableFuture.globalExceptionHandler = handler;
    }

    /**
     * 创建一个新的 MyCompletableFuture 实例并执行异步任务
     */
    public static <U> MyCompletableFuture<U> supplyAsync(Supplier<U> supplier) {
        return new MyCompletableFuture<>().supplyAsyncCustom(supplier);
    }

    /**
     * 创建一个新的 MyCompletableFuture 实例并执行异步任务
     */
    public static MyCompletableFuture<Void> runAsync(Runnable runnable) {
        return new MyCompletableFuture<>().runAsyncCustom(runnable);
    }

    /**
     * 创建一个新的 MyCompletableFuture 实例并执行异步任务，指定线程池
     */
    public static <U> MyCompletableFuture<U> supplyAsync(Supplier<U> supplier, Executor executor) {
        if (executor == null) {
            executor = ThreadPoolHelper.asyncTaskThreadPool;
        }
        return new MyCompletableFuture<>().supplyAsyncCustom(supplier, executor);
    }

    /**
     * 创建一个新的 MyCompletableFuture 实例并执行异步任务，指定线程池
     */
    public static MyCompletableFuture<Void> runAsync(Runnable runnable, Executor executor) {
        if (executor == null) {
            executor = ThreadPoolHelper.asyncTaskThreadPool;
        }
        return new MyCompletableFuture<>().runAsyncCustom(runnable, executor);
    }

    /**
     * 重写 handle 方法，捕获异常并传递给全局异常处理器
     */
    @Override
    public <U> MyCompletableFuture<U> handle(BiFunction<? super T, Throwable, ? extends U> fn) {
        // 创建一个新的 MyCompletableFuture 实例
        MyCompletableFuture<U> newFuture = new MyCompletableFuture<>();
        // 调用父类的 handle 方法，并将结果传递给新的 MyCompletableFuture
        super.handle((result, ex) -> {
            if (ex != null && globalExceptionHandler != null) {
                // 调用全局异常处理器
                globalExceptionHandler.accept(ex, "Exception in MyCompletableFuture.handle()");
            }
            return fn.apply(result, ex);
        }).whenComplete((value, exception) -> {
            if (exception == null) {
                // 如果没有异常，完成新的 MyCompletableFuture
                newFuture.complete(value);
            } else {
                // 如果有异常，完成新的 MyCompletableFuture 并传递异常
                newFuture.completeExceptionally(exception);
            }
        });
        return newFuture;
    }

    /**
     * 重写 whenComplete 方法，捕获异常并传递给全局异常处理器
     */
    @Override
    public MyCompletableFuture<T> whenComplete(BiConsumer<? super T, ? super Throwable> action) {
        return (MyCompletableFuture<T>) super.whenComplete((result, ex) -> {
            if (ex != null && globalExceptionHandler != null) {
                // 调用全局异常处理器
                globalExceptionHandler.accept(ex, "Exception in MyCompletableFuture");
            }
            action.accept(result, ex);
        });
    }

    /**
     * 重写 exceptionally 方法，捕获异常并传递给全局异常处理器
     */
    @Override
    public MyCompletableFuture<T> exceptionally(Function<Throwable, ? extends T> fn) {
        return (MyCompletableFuture<T>) super.exceptionally(ex -> {
            if (ex != null && globalExceptionHandler != null) {
                // 调用全局异常处理器
                globalExceptionHandler.accept(ex, "Exception in MyCompletableFuture.exceptionally()");
            }
            return fn.apply(ex);
        });
    }

    /**
     * 重写 join 方法，捕获异常并传递给全局异常处理器
     */
    @Override
    public T join() {
        try {
            return super.join();
        } catch (CompletionException e) {
            if (e.getCause() != null && globalExceptionHandler != null) {
                // 调用全局异常处理器
                globalExceptionHandler.accept(e.getCause(), "Exception in MyCompletableFuture.join()");
            }
            throw e;
        }
    }

    /**
     * 重写 get 方法，捕获异常并传递给全局异常处理器
     */
    @Override
    public T get() throws InterruptedException, ExecutionException {
        try {
            return super.get();
        } catch (ExecutionException e) {
            if (e.getCause() != null && globalExceptionHandler != null) {
                // 调用全局异常处理器
                globalExceptionHandler.accept(e.getCause(), "Exception in MyCompletableFuture.get()");
            }
            throw e;
        }
    }

    /**
     * 重写 completeExceptionally 方法，捕获异常并传递给全局异常处理器
     */
    @Override
    public boolean completeExceptionally(Throwable ex) {
        if (ex != null && globalExceptionHandler != null) {
            // 调用全局异常处理器
            globalExceptionHandler.accept(ex, "Exception in MyCompletableFuture.completeExceptionally()");
        }
        return super.completeExceptionally(ex);
    }

    /**
     * 重写 supplyAsync 方法，确保所有异步任务都使用 MyCompletableFuture
     */
    public <U> MyCompletableFuture<U> supplyAsyncCustom(Supplier<U> supplier) {
        MyCompletableFuture<U> future = new MyCompletableFuture<>();
        CompletableFuture.supplyAsync(() -> {
            try {
                return supplier.get();
            } catch (Throwable ex) {
                if (globalExceptionHandler != null) {
                    // 调用全局异常处理器
                    globalExceptionHandler.accept(ex, "Exception in MyCompletableFuture.supplyAsync()");
                } else {
                    // 如果没有设置异常处理器，则使用默认的异常处理器打印异常堆栈信息
                    MyCompletableFuture.setGlobalExceptionHandler((e, message) -> {
                        log.error("线程池CompletableFuture.supplyAsync(Supplier<U> supplier)执行任务异常：{}", message, e);
                    });
                }
                future.completeExceptionally(ex);
                throw ex;
            }
        }, ThreadPoolHelper.asyncTaskThreadPool).whenComplete((result, ex) -> {
            if (ex == null) {
                future.complete(result);
            } else {
                future.completeExceptionally(ex);
            }
        });
        return future;
    }

    /**
     * 重写 runAsync 方法，确保所有异步任务都使用 MyCompletableFuture
     */
    public MyCompletableFuture<Void> runAsyncCustom(Runnable runnable) {
        MyCompletableFuture<Void> future = new MyCompletableFuture<>();
        CompletableFuture.runAsync(() -> {
            try {
                runnable.run();
            } catch (Throwable ex) {
                if (globalExceptionHandler != null) {
                    // 调用全局异常处理器
                    globalExceptionHandler.accept(ex, "Exception in MyCompletableFuture.runAsync()");
                } else {
                    // 如果没有设置异常处理器，则使用默认的异常处理器打印异常堆栈信息
                    MyCompletableFuture.setGlobalExceptionHandler((e, message) -> {
                        log.error("线程池CompletableFuture.runAsync(Runnable runnable)执行任务异常：{}", message, e);
                    });
                }
                future.completeExceptionally(ex);
                throw ex;
            }
        }, ThreadPoolHelper.asyncTaskThreadPool).whenComplete((result, ex) -> {
            if (ex == null) {
                future.complete(null);
            } else {
                future.completeExceptionally(ex);
            }
        });
        return future;
    }

    /**
     * 重写 supplyAsync 方法，指定线程池，确保所有异步任务都使用 MyCompletableFuture
     */
    public <U> MyCompletableFuture<U> supplyAsyncCustom(Supplier<U> supplier, Executor executor) {
        MyCompletableFuture<U> future = new MyCompletableFuture<>();
        CompletableFuture.supplyAsync(() -> {
            try {
                return supplier.get();
            } catch (Throwable ex) {
                if (globalExceptionHandler != null) {
                    // 调用全局异常处理器
                    globalExceptionHandler.accept(ex, "Exception in MyCompletableFuture.supplyAsync()");
                } else {
                    // 如果没有设置异常处理器，则使用默认的异常处理器打印异常堆栈信息
                    MyCompletableFuture.setGlobalExceptionHandler((e, message) -> {
                        log.error("线程池CompletableFuture.supplyAsync(Supplier<U> supplier, Executor executor)执行任务异常：{}", message, e);
                    });
                }
                future.completeExceptionally(ex);
                throw ex;
            }
        }, executor).whenComplete((result, ex) -> {
            if (ex == null) {
                future.complete(result);
            } else {
                future.completeExceptionally(ex);
            }
        });
        return future;
    }

    /**
     * 重写 runAsync 方法，指定线程池，确保所有异步任务都使用 MyCompletableFuture
     */
    public MyCompletableFuture<Void> runAsyncCustom(Runnable runnable, Executor executor) {
        MyCompletableFuture<Void> future = new MyCompletableFuture<>();
        CompletableFuture.runAsync(() -> {
            try {
                runnable.run();
            } catch (Throwable ex) {
                if (globalExceptionHandler != null) {
                    // 调用全局异常处理器
                    globalExceptionHandler.accept(ex, "Exception in MyCompletableFuture.runAsync()");
                } else {
                    // 如果没有设置异常处理器，则使用默认的异常处理器打印异常堆栈信息
                    MyCompletableFuture.setGlobalExceptionHandler((e, message) -> {
                        log.error("线程池CompletableFuture.runAsync(Runnable runnable, Executor executor)执行任务异常：{}", message, e);
                    });
                }
                future.completeExceptionally(ex);
                throw ex;
            }
        }, executor).whenComplete((result, ex) -> {
            if (ex == null) {
                future.complete(null);
            } else {
                future.completeExceptionally(ex);
            }
        });
        return future;
    }
}